#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>
#include <string.h>
#include <gmp.h>
#include <unistd.h>
#include <sys/types.h>
#include "mpzqform.h"
#include "mpzsmoothscan.h"
#include "smoothrelation.h"
/*#include "suborder_relations.hpp"*/

/*using namespace ANTL;
using namespace std;*/

/*
    Copyright 2009-2010 Andrew V. Sutherland

    This file is part of "smoothrelation".

    smoothrelation is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 2 of the License, or
    (at your option) any later version.

    smoothrelation is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with smalljac.  If not, see <http://www.gnu.org/licenses/>.
*/





int test_drew ( int argc, char *argv[] , unsigned long input_seed)
{
	clock_t nstart, start, end;
	gmp_randstate_t rstate;
	int e[MAX_K0], s[MAX_K0], etot, etot2, esubtot;
	double maxNwt, wt, Bwt;
	struct exp_struct in_exp, out_exp;
	mpz_t u,v,x,y,p,q,D,L,M,out_a, out_b,d,maxN, P, Pa, Pb;
	unsigned long B, E,T, n, N, B0, tests, q0;
	register int i, j, k0, k, lastj, m, m0;
	double z;
	char sign;
	int cnt, badcnt, btests, arity, unramified, i0;
        int ret_int;

	if ( argc < 3 ) {
		puts ("smoothrelation D B [P unramified seed k0]");
		puts ("    D is an imaginary discriminant (if positive, it will be made negative)");
		puts ("    B is a smoothness bound");
		puts ("    P is an optional split prime that will be included in the relation.");
		puts ("    unramified is a flag indicating that only unramified primes should be used in the factor base.");
		puts ("    seed is a positive integer seed for generating random numbers.");
		return 0;
	}

	mpz_init(p);  mpz_init(q); mpz_init(D); mpz_init(L); mpz_init(M); mpz_init(d); mpz_init(u); mpz_init(v); mpz_init(x); mpz_init(y); mpz_init(maxN); mpz_init(out_a); mpz_init(out_b); mpz_init(P); mpz_init(Pa); mpz_init(Pb);
	mpz_set_str(D,argv[1],0);
	if ( mpz_sgn(D)>0 ) mpz_neg(D,D);
	if ( ! mpz_qform_valid_discriminant(D) ) { gmp_printf ("D=%Zd is not a valid discriminant\n", D); return 0; }
	mpz_neg(L,D);
	mpz_fdiv_q_ui(maxN,L,3);									// bound on largest possible norm in cl(D), sqrt(-D/3)
	mpz_sqrt(maxN,maxN);  mpz_add_ui(maxN,maxN,1);	
	maxNwt = mpz_sizeinbase(maxN,2);						// upper bound on bits in largest possible norm
	mpz_root(L,L,4);
	B = atol (argv[2]);
	if ( argc > 3 ) {
		mpz_set_str(P,argv[3],0);
		if ( mpz_sgn(P) )
			if ( ! mpz_qform_primeform2(Pa,Pb,P,D) ) { gmp_printf ("Couldn't create primeform for P=%Zd\n", P); return 0; }
	}
	if ( argc > 4 ) unramified = atoi(argv[4]); else unramified = 0;
	gmp_randinit_default (rstate);
	if ( argc > 5 ) { 
		N = atol (argv[5]);
	} else {
		N = (((unsigned long)gethostid())<<32) + getpid();		// use a seed that depends on process id to ensure parallel jobs generate independent random forms
		N *= time(0);										// and throw in the current time for good measure
	}
	if ( argc > 6 ) k0 = atoi(argv[6]); else k0 = 4;
	if ( k0 < 3 || k0 > MAX_K0 ) { printf ("Invalid k0=%d, must be between 2 and %d\n", k0, MAX_K0); return 0; }
//	k0 = MAX_K0;											// currently k0 is hardwired to MAX_K0, typically set to 3 or 4
	E = 1UL<<40;
        N = input_seed;
        // we set E at an arbitrarily large value that we never plan to reach, relying on the user to pick a good B
	gmp_randseed_ui (rstate, N);
	printf ("Using seed %lu\n", N);
	for ( i = 0 ; i < MAX_BASE ; i++ ) { mpz_init(a[i]); mpz_init(b[i]); }

	gmp_printf ("Searching for relation in cl(%Zd)\n", D);
	
restart:
	
	start = clock();
	tests = 0;
	cnt = 0;
	for (;;) {
		m = 0;
		mpz_set_ui(p,2);
		mpz_set_ui(M,1);
		smoothbox_init(D,B,unramified);
		while ( m < MAX_BASE && mpz_cmp_ui(p,B) < 0 ) {
			mpz_mul(x,p,p);
			if ( ! mpz_divisible_p (D,x) && mpz_kronecker(D,p) > (unramified ? 0 : -1) ) {
				ell[m++] = mpz_get_ui(p);
				for ( mpz_set(q,p) ; mpz_cmp(q,maxN) <= 0 ; mpz_mul(q,q,p) );
//				mpz_mul(M,M,q);
			}
			mpz_nextprime(p,p);
		}
		printf ("Factor base has %d elements, largest norm is %ld\n", m, ell[m-1]);
//		printf ("M has %d bits\n", mpz_sizeinbase(M,2));
		
		// Move norms of ideal classes with order 2 to the end of the list
		for ( i = k=0 ; i < m ; i++ ) {
			if ( mpz_divisible_ui_p(D,ell[i]) ) {
				mpz_set_ui(p,ell[i]);
				if ( ! mpz_qform_primeform2(x,y,p,D) ) { printf ("Fatal error, couldn't create primeform for ell=%ld\n",ell[i]); return 0; }
				mpz_set(v,y);
				mpz_qform_invert2(x,y,D);
				if ( mpz_cmp(v,y)==0 ) { ell[m+k++] = ell[i];  ell[i] = 0; }
			}
		}
		for ( i = j = 0 ; i < m+k ; i++ ) if ( ell[i] ) ell[j++] = ell[i];
		m0 = m-k;
		
		for ( i = 0 ; i < m ; i++ ) {
			mpz_set_ui(p,ell[i]);
			if ( ! mpz_qform_primeform2(a[i],b[i],p,D) ) { printf ("Fatal error, couldn't create primeform for ell=%ld\n",ell[i]); return 0; }
		}

#if TEST_TIMING
//		Timing code to test the speed of the group operation and smoothness testing
		mpz_set(u,a[m-1]); mpz_set(v,b[m-1]);
		nstart = clock();
		for ( i = 0 ; i < 1000000 ; i++ ) mpz_qform_compose2(u,v,u,v,a[m-2],b[m-2],D,L,1);
		end = clock();
		printf ("Time to compose 1000000 forms is %ld msecs (%ld/sec)\n", delta_msecs(nstart,end), 1000000000/(delta_msecs(nstart,end)+1));

		mpz_set(u,a[m-1]); mpz_set(v,b[m-1]);
		k = 0;
		nstart = clock();
		for ( i = 0 ; i < 1000000 ; i++ ) {
			mpz_qform_compose2(u,v,u,v,a[m-2],b[m-2],D,L,1);
			if ( smoothbox_feed(x,y,&out_exp,u,v,&in_exp) ) k++;
		}
		end = clock();
		printf ("Time to enum+test 1000000 forms for %ld-smoothness using Bernstein %ld msecs (%ld/sec) (found %d)\n", B, delta_msecs(nstart,end), 1000000000/(delta_msecs(nstart,end)+1), k);
#endif
		
		// compute exponent bounds based on B
//printf ("Exponent bounds for B=%d: ", B);
		Bwt = log2_Drew(B)+1;
		for ( i = j = 0, etot = 0 ; i < m0 ; i++ ) {
			wt = log2_Drew(ell[i]);
			ebits[i] = (int)ceil(OMEGA*(Bwt-wt));
			if ( ebits[i] > MAX_EXP_BITS ) ebits[i] = MAX_EXP_BITS;
			eoff[i] = etot;
//printf ("%d: %d bits at offset %d\n", ell[i],ebits[i], eoff[i]);
			if ( etot+ebits[i] > MAX_EBITS ) { printf ("MAX_EBITS=%d is too small, increase to at least %d\n", MAX_EBITS, etot); return 0; }
			mpz_init_set(ga[j],a[i]); mpz_init_set(gb[j],b[i]);
			bitmap[j] = i;
			wts[j]=wt;
			for ( j++ ; j < etot + ebits[i] ; j++ ) {
				mpz_qform_square2(ga[j],gb[j],ga[j-1],gb[j-1],D,L);
				bitmap[j] = i;
				wts[j] = (j-eoff[i])*wt;
//printf ("Bit %d has wt %.2f\n", j, wts[j]);
			}
			etot+=ebits[i];
		}
//printf("Exponent bits = %d, Max norm bits = %d\n", etot,(int)maxNwt);
		for ( i = etot2 = 0 ; i < m0 ; i++ ) { if ( ebits[i] < 2 ) break; etot2 += ebits[i]; }
		for ( N = 0 ; N < E ; N+=BATCH_SIZE ) {
			// pick starting point by setting k0 random bits in beta's exponent vector (making sure ell[s[i]] we use are distinct)
			memset(beta_e,0,etot);
			esubtot = 0;
			Bwt = 0.0;
			for ( i = 0 ; i < k0 ; i++ ) {
				do {
					j = gmp_urandomm_ui (rstate,etot2);
					s[i] = bitmap[j];
					for ( k = 0 ; k < i ; k++ ) if ( s[k]==s[i] ) break;
				} while ( k < i );
				esubtot += ebits[s[i]];
//printf("j=%d, s[%d]=%d\n", j, i, ell[s[i]]);
				Bwt += wts[j];
				beta_e[j] = 1;
				in_exp.e[i] = (1<<(j-eoff[s[i]]));
				in_exp.i[i] = s[i];
				if ( i ) {
					mpz_qform_compose2(u,v,u,v,ga[j],gb[j],D,L,1);
				} else {
					if ( ! mpz_sgn(Pa) ) {
						mpz_set(u,ga[j]); mpz_set(v,gb[j]);
					} else {
						mpz_qform_compose2(u,v,Pa,Pb,ga[j],gb[j],D,L,1);
					}
				}
			}
			// Make sure we have enough bits to work with
			if ( esubtot < LOG2_BATCH_SIZE ) { N-=BATCH_SIZE; continue; }
			badcnt = 0;
			lastj = -1;
			btests = 0;
			for ( n = 0 ; n < BATCH_SIZE ; n++ ) {
				if ( Bwt < maxNwt ) goto next_beta;
				badcnt = 0;
				tests++; btests++;
//				if ( mpz_divisible_p(M,u) ) {					// Brute force divisibility test for smoothness, no longer used
				if ( smoothbox_feed (out_a,out_b,&out_exp,u,v,&in_exp) ) {
//gmp_printf("Smooth form <%Zd,%Zd,*>\n", out_a, out_b);
					// now factor the norm with trial division, which is easy since it is smooth
					mpz_set(x,out_a);
					for ( i = 0 ; i < m ; i++ ) {
						for ( h[i] = 0 ; mpz_divisible_ui_p(x,ell[i]) ; h[i]++ ) mpz_divexact_ui(x,x,ell[i]);
						if ( h[i] ) hs[i] = ( mpz_fdiv_ui(out_b,2*ell[i]) > ell[i] ? -1 : 1 ); else hs[i] = 0;
					}
					// determine exponents for k0 primeforms; if these intersect the factorization of a, we need to add/sub, possibly changing signs
					// our sign convention will be to move the k0 forms to the other side, negating their signs (which were all nonnegative to begin with)
					for ( i0 = 0 ; i0 < k0 ; i0++ ) {
						i = out_exp.i[i0];  j = out_exp.e[i0];
						if ( h[i] ) {
							if ( hs[i] < 0 ) { h[i] += j; } else { if ( h[i] > j ) { h[i] -= j; } else { h[i] = j - h[i]; hs[i] = -1; } };
						} else {
							h[i] = j;  hs[i] = -1;
						}
					}
					// make sure relation is really new (only relevant when searching for multiple relations)
					for ( i = 0 ; i < cnt ; i++ ) {
						for ( j = 0 ; j < m ; j++ ) if ( h[j] != r[i][j] ) break;
						if ( j == m ) break;
					}
					if ( i == cnt ) {
						for ( j = 0, arity = 0 ; j < m ; j++ ) { r[i][j] = h[j]; if ( h[j] ) arity++; }
						if ( mpz_sgn(P) ) { 
							gmp_printf ("Smooth relation for %Zd = \n    ", P);
						} else {						
							printf ("Smooth %d-relation: ", arity);
						}
                                                ret_int = (int)m;
						for ( i = 0 ; i < m ; i++ ) if ( h[i] ) printf ("%ld^%s%ld ", ell[i], (hs[i]<0?"-":""), h[i]);
						end = clock();
						cnt++;
						if ( cnt < MAX_RELATIONS ) {
							printf("    %ld msecs\n", delta_msecs(start,end));
							printf ("%ld tests total, %ld tests/s\n", tests, 1000UL*tests/(delta_msecs(start,end)+1));
						} else {
							puts("");
							goto done;
						}
					}
					break;
				}
next_beta:
//printf("at next beta: Bwt=%f, maxNwt=%f\n", Bwt, maxNwt);
				if ( badcnt++ > 10 ) break;
				do { j = gmp_urandomm_ui(rstate,esubtot); } while ( j==lastj);
				lastj = j;
				for ( i = 0 ; j >= ebits[s[i]] ; i++ )  j -= ebits[s[i]];
				j += eoff[s[i]];
				if ( beta_e[j] ) {
					if ( Bwt - wts[j] < maxNwt ) goto next_beta;
//printf ("clearing bit %d (%d^%d)\n", j, ell[s[i]], (1<<(j-eoff[s[i]])));
					mpz_set(x,gb[j]);
					mpz_qform_invert2(ga[j],x,D);
					mpz_qform_compose2(u,v,u,v,ga[j],x,D,L,1);
					Bwt -= wts[j]; 
					in_exp.e[i] ^= (1<<(j-eoff[s[i]]));
				} else {
//printf ("setting bit %d (%d^%d)\n", j, ell[s[i]], (1<<(j-eoff[s[i]])));
					mpz_qform_compose2(u,v,u,v,ga[j],gb[j],D,L,1);
					Bwt += wts[j];
					in_exp.e[i] |= (1<<(j-eoff[s[i]]));
				}
				beta_e[j]=1-beta_e[j];
				if ( Bwt < maxNwt ) goto next_beta;
			}
//printf ("%d of %d in batch tested\n", btests, BATCH_SIZE);
		}
		// currently we never get here because E is hard-wired to a ridiculously big value
		if ( 1.01*B < MAX_B ) {
			B = 1.05*B+1;
			E = 1.05*1.05*1.05*E+1;
			printf ("B=%ld, E=%ld\n", B, E);
		} else {
			printf ("At MAX_B=%d\n", MAX_B);
		}
	}
		

done:
	end = clock();
	T = delta_msecs(start,end);
	printf ("Performed a total of %ld tests, %ld tests/sec, total elapsed time %.3f secs\n", tests, (T?(1000UL*tests)/T:-1), (double)T/1000.0);
        return ret_int;
}


/*
	smoothbox should eventually be put into its own module
*/


void smoothbox_init (mpz_t D, int B, int unramified)
{
	static mpz_t p,q;
	int i;
	
	if ( ! sb_a ) {
		mpz_init(p);  mpz_init(q);
		sb_exp = malloc(SB_BS*sizeof(struct exp_struct));
		sb_a = malloc(SB_BS*sizeof(mpz_t));
		for ( i = 0 ; i < SB_BS ; i++ ) mpz_init(sb_a[i]);
		sb_b = malloc(SB_BS*sizeof(mpz_t));
		for ( i = 0 ; i < SB_BS ; i++ ) mpz_init(sb_b[i]);
		sb_t = malloc(2*SB_BS*sizeof(mpz_t));
		for ( i = 0 ; i < 2*SB_BS ; i++ ) mpz_init(sb_t[i]);
		mpz_init(sb_P);
	}
	// don't bother being efficient about computing sb_P, we expect B to be fairly small
	mpz_set_ui(p,1);
	sb_nsp = 0;
	mpz_set_ui(sb_P,1);
	for (;;) {
		do { mpz_nextprime(p,p); mpz_mul(q,p,p); } while ( mpz_kronecker(D,p) < (unramified ? 1 : 0)  || mpz_divisible_p(D,q) );		// AVS 01/28/12 changed to handle ramified flag and to not use primes that divide the conductor
		if ( mpz_cmp_ui(p,B) > 0 ) break;
		if ( sb_nsp < SB_NSP ) { sb_sp[sb_nsp] = mpz_get_ui(p); sb_spw[sb_nsp] = log2_Drew(sb_sp[sb_nsp]);  sb_nsp++; }
		for ( mpz_set(q,p) ; mpz_cmp_ui(q,B*B) <= 0 ; mpz_mul(q,q,p), mpz_mul(sb_P,sb_P,p) );
	}
	sb_n = 0;
}

int smoothbox_feed (mpz_t out_a, mpz_t out_b, struct exp_struct out_exp[1] , mpz_t a, mpz_t b, struct exp_struct in_exp[1])
{
	register double wt;
	register unsigned long q;
	register int i;

	i = ( sb_nsp && sb_sp[0] == 2 ? 1 : 0 );
	wt = ( i ? mpz_scan1(a,0) : 0.0);
	for ( ; i < sb_nsp ; i++ )
		for ( q = sb_sp[i] ; mpz_divisible_ui_p(a,q) ; q *= sb_sp[i] ) wt += sb_spw[i];
	if ( wt >= SB_MIN_WT ) {
		mpz_set(sb_a[sb_n],a);  mpz_set(sb_b[sb_n],b); sb_exp[sb_n]=in_exp[0]; sb_n++;
		if ( sb_n==SB_BS ) return smoothbox_flush(out_a, out_b, out_exp);
	}
	return 0;
}

int smoothbox_flush (mpz_t out_a, mpz_t out_b, struct exp_struct out_exp[1])
{
	int i;

	for ( i = sb_n ; i < SB_BS ; i++ ) mpz_set_ui(sb_a[i],1);
	i = mpz_smoothscan (sb_a, SB_BS, 0, sb_P, sb_t);
	if ( i < sb_n ) {
		mpz_set(out_a, sb_a[i]);  mpz_set(out_b, sb_b[i]);  out_exp[0] = sb_exp[i];
		sb_n = 0; return 1;
	}
	sb_n = 0;  return 0;
}

